from .Matrix import Matrix
from .Vector import Vector

class LinearSystem:
    def __init__(self,A,b):
        assert A.row_num() == len(b), "row number of A must be eual to the length"

        self._m = A.row_num()
        self._n = A.col_num()

        # 这里Ab是列表
        self.Ab = [
            Vector(A.row_vector(i)[:] + [b[i]])
            for i in range(self._m)
        ]

        pass
        
    def gauss_jordan_elim(self):
        self._forward()
        self._backward()

        pass

    def _forward(self):
        '''从上到下将对角线左侧清零'''
        n = self._m

        for i in range(n):
            #Ab[i][i]为主元
            max_row = self._max_row(i,n)

            self.Ab[i],self.Ab[max_row] = self.Ab[max_row],self.Ab[i]

            #归一
            self.Ab[i] = self.Ab[i] / self.Ab[i][i]

            for j in range(i+1,n):
                self.Ab[j] = self.Ab[j] - self.Ab[i]*self.Ab[j][i]

        pass

    def _max_row(self,index,n):
        best,ret = self.Ab[index][index],index

        for i in range(index + 1,n):
            if self.Ab[i][index]<best:
                best,ret = self.Ab[i][index],i

        return ret

    def _backward(self):
        '''从下到上将对角线右侧清零'''
        n = self._m
        
        #三个参数时是[start,stop) step
        for i in range(n-1,0,-1):
            for j in range(i - 1,-1,-1):             
                self.Ab[j] = self.Ab[j] - self.Ab[i]*self.Ab[j][i]

        pass

    def fancy_print(self):
        for i in range(self._m):
            print(" ".join(str(self.Ab[i][j]) for j in range(self._n)),end=" ")
            print("|",self.Ab[i][-1])
        
